import numpy as np
import pandas as pd

df = pd.DataFrame({
    "quarter": [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4],
    "x": np.random.randn(12),
    "date": pd.date_range("2024-01-01", periods=12, freq="MS")
})

df.head()

df.x

df.columns

df.index

accidents = pd.read_csv("data/estonia-traffic-accidents-clean.csv")

accidents.shape

accidents.columns

accidents.head()

accidents["date"].head(4)

accidents[["date", "county"]].head(4)

accidents["date"] = pd.to_datetime(accidents["date"])
type(accidents["date"][0])

start = accidents["date"].min()
end = accidents["date"].max()
print(f"First accident: {start}\nLast accident: {end}")

accidents["accident_type"].value_counts()

accidents["county"] == "Harju maakond"

accidents["killed_or_injured"] = accidents["killed"] + accidents["injured"]
accidents[['killed', 'injured', 'killed_or_injured']].head()

accidents["killed_or_injured"].sum()

by_type = accidents.groupby("accident_type").agg({"killed_or_injured": "sum"})

by_type

accidents.select(
    # select 'killed'
    pl.col("killed")
    # add 'injured'
    .add(pl.col("injured"))
    # give the result a new column name
    .alias("killed_or_injured")
    # identify the max value
    .max())

accidents = accidents.with_columns(
    pl.col("killed").add(pl.col("injured")).alias("killed_or_injured"),
    pl.col("killed").add(pl.col("injured")).truediv(pl.col("persons_involved")).alias("harmed_rate")
)
accidents.select(["date", "persons_involved", "killed_or_injured", "harmed_rate"]).head(5)

by_county = (accidents
             .group_by("county")
             .agg(pl.col("killed_or_injured").sum())
             .sort("killed_or_injured", descending=True)
)
by_county.head()

county_pop = (pl.read_csv("data/county-pop.csv", skip_rows=2)
              .rename({"County": "county", "Age groups total": "population"})
              .select(["county", "population"])
              # this has "county" in the county names, not "maakond"
              .with_columns(pl.col("county").str.replace("county", "maakond"))
              )

by_county_w_pop = by_county.join(county_pop, on="county", how="left")
by_county_w_pop.head(3)

by_county_w_pop.select(
    pl.col("county"), 
    pl.col("killed_or_injured"),
    pl.col("killed_or_injured").truediv(pl.col("population")).mul(1000).alias("rate/1000")
    ).head(3)

reflectors = (pl.read_csv("data/reflectors.csv", has_header=True, separator=";", skip_rows=2)
              .filter(pl.col("Sex").ne("Men and women"))
              .drop(["Type of data", "Year", "All age groups (16-64)"])
)
reflectors.head()

(reflectors
 .with_columns(pl.col("Reflector use").str.replace("Never walk on dark streets, roads", "Never"))
 .group_by(["Reflector use", "Sex", "age_group"])
 .agg(pl.col("percentage").sum())
 .filter(pl.col("Reflector use").eq("Never"))
 .sort(["age_group", "Sex"])
 .plot.line(x = "age_group", y = "percentage", color = "Sex")
 .properties(width=700, height=300)
)